import os
import torchaudio
import torch
import random
from torch.utils.data import Dataset

random.seed(0)


def slice_audio(length: int, audio: torch.Tensor):
    if audio.size(1) > length:
        return audio[:, :length]
    else:
        audio_0 = torch.zeros(1, length)
        audio_0[:, : audio.size(1)] = audio
        return audio_0


class SpexData(Dataset):
    def __init__(self, tr_dir, sr_dir):
        tr_dir_mix = tr_dir + "mix"
        self.s1_dir = tr_dir + "s1"
        self.s2_dir = tr_dir + "s2"
        self.sr_dir = sr_dir
        self.sr = os.listdir(sr_dir)
        self.sr_len = len(self.sr)
        self.tr_mix = [
            tr_dir_mix + "/" + f for f in os.listdir(tr_dir_mix) if f.endswith(".wav")
        ]
        self.half_len = len(self.tr_mix)

    def __len__(self):
        return self.half_len

    def __getitem__(self, index):
        ## TODO: disable second speaker
        second_speaker = False
        mix = self.tr_mix[index]
        filename = os.path.split(mix)[1]
        ## output
        if second_speaker:
            speaker = (filename.split("_")[2])[:3]
            output = self.s2_dir + "/" + os.listdir(self.s2_dir)[index]
        else:
            speaker = (filename.split("_")[0])[:3]
            output = self.s1_dir + "/" + os.listdir(self.s1_dir)[index]
        ## speaker encoder
        speaker_directory = self.sr_dir + "/" + speaker
        files = [
            speaker_directory + "/" + f
            for f in os.listdir(speaker_directory)
            if f.endswith(".wav")
        ]
        index = random.randint(0, len(files) - 1)
        encoder = files[index]

        ## cut the audio
        # 4 * 8000 = 32000
        mix: torch.Tensor = torchaudio.load(mix)[0]
        encoder: torch.Tensor = torchaudio.load(encoder)[0]
        output: torch.Tensor = torchaudio.load(output)[0]
        mix = slice_audio(32000, mix)
        encoder = slice_audio(32000, encoder)
        output = slice_audio(32000, output)

        embedding = torch.zeros(self.sr_len)
        embedding[self.sr.index(speaker)] = 1

        return mix, encoder, output, embedding


if __name__ == "__main__":
    pass
